import os
os.environ['DISABLE_ARGUMENTS_CHECKER'] = "Yes"

import talib
from rqalpha.apis import *

# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def log_cash(context, bar_dict):
    logger.info("Remaning cash: %r" % context.portfolio.cash)
    pass


# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):
    # 全部加载后才有这些模块接口
    from rqalpha.api import get_future_contracts, scheduler, market_open

    context.s1 = context.config.params.universe

    # 设置这个策略当中会用到的参数，在策略中可以随时调用，这个策略使用长短均线，我们在这里设定长线和短线的区间，在调试寻找最佳区间的时候只需要在这里进行数值改动
    context.SHORTPERIOD = 20
    context.LONGPERIOD = 60
    # windows
    context.window = context.LONGPERIOD + 1
    
    # 初始化时订阅合约行情。订阅之后的合约行情会在handle_bar中进行更新
    subscribe(context.s1)
    # fc = get_future_contracts(context.s1[0:2])
    # print(fc)

    # scheduler.run_daily(log_cash, time_rule=market_open(minute=10))
    scheduler.run_monthly(log_cash, tradingday=1)
    pass


# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def handle_bar(context, bar_dict): # 1.9秒
    # 开始编写你的主要的算法逻辑
    # bar_dict[order_book_id] 可以拿到某个证券的bar信息
    # context.portfolio 可以拿到现在的投资组合状态信息
    
    # 因为策略需要用到均线，所以需要读取历史数据
    prices = history_bars(context.s1, context.window, '1m', 'close', adjust_type="pre")
    # 品种最开始周期,需要跳过剔除
    if len(prices) < context.window:
      return
    # 使用talib计算长短两根均线，均线以array的格式表达
    short_avg = talib.SMA(prices, context.SHORTPERIOD)
    long_avg = talib.SMA(prices, context.LONGPERIOD)

    # 获取当前投资组合中股票的仓位
    long_pos_a = get_position(context.s1, POSITION_DIRECTION.LONG)
    short_pos_a = get_position(context.s1, POSITION_DIRECTION.SHORT)

    # 如果短均线从上往下跌破长均线，也就是在目前的bar短线平均值低于长线平均值，而上一个bar的短线平均值高于长线平均值
    if short_avg[-1] - long_avg[-1] < 0 and short_avg[-2] - long_avg[-2] > 0:
        if long_pos_a.quantity > 0:
          sell_close(context.s1, long_pos_a.quantity)
          logger.info('买入仓位平仓成功')
        if short_pos_a.quantity <= 0:
          sell_open(context.s1, 1)
          logger.info('卖出仓位创建成功')

    # 如果短均线从下往上突破长均线，为入场信号
    if short_avg[-1] - long_avg[-1] > 0 and short_avg[-2] - long_avg[-2] < 0:
        if short_pos_a.quantity > 0:
          buy_close(context.s1, short_pos_a.quantity)
          logger.info('卖出仓位平仓成功')
        if long_pos_a.quantity <= 0:
          buy_open(context.s1, 1)      
          logger.info('买入仓位创建成功')

    if context.now.date() == context.run_info.end_date:
      if context.now.hour == 15:
        # 获取当前投资组合中股票的仓位，年底清仓
        long_pos_a = get_position(context.s1, POSITION_DIRECTION.LONG)
        short_pos_a = get_position(context.s1, POSITION_DIRECTION.SHORT)
        if short_pos_a.quantity > 0:
          buy_close(context.s1, short_pos_a.quantity)
          logger.info('清仓成功')
        if long_pos_a.quantity > 0:
          sell_close(context.s1, long_pos_a.quantity)
          logger.info('清仓成功')
    pass


__config__ = {
  "params": {
    "universe": "IH88"
  },
  "base": {
    "start_date": "2012-01-01",
    "end_date": "2021-12-31",
    "frequency": "1m",
    "accounts": {
        "future": 1000000
    }
  },
  "extra": {
    "log_level": "error",
  },
  "mod": {
    "sys_analyser": {
      "benchmark": "000300.XSHG",
      "enabled": True,
      "plot": False
    },
    "sys_progress": {
      "enabled": False,
      "show": False
    },
    "sys_risk": {
      "enabled": True,
      "validate_price": True
    },
    "sys_simulation": {
      "enabled": True,
      "signal": False
    },
    "sys_transaction_cost": {
      "enabled": True
    }
  }
}


if __name__ == '__main__':
    # 您可以指定您要传递的参数
    import time
    import json
    # from rqalpha import run_func
    # import cProfile
    # pr = cProfile.Profile()
    # pr.enable()
    # start = time.time()
    # ret = run_func(init=init, handle_bar=handle_bar, config=__config__)
    # print(time.time() - start, ret)
    # pr.disable()
    # pr.print_stats('cumtime')
    # 批量回测
    from nextt.tasks import year_split_file_tasks

    code_list = \
      ['SN','CY','SM','RM','RU','LH','CF','AG','BC','JM','RS','PM','IF','I','SC','A','WH','MA','T','TF','CJ','FB','BU','SS','CS','EG','LU','V','AP','UR',\
        'NR','L','J','HC','SF','RR','FU','SP','TS','PP','CU','WR','BB','PB','C','P','JR','IH','JD','ZN','LR','PG','M','EB','PK','IC','ZC','AL','FG','OI',\
        'RI','SA','B','Y','RB','SR','AU','TA','NI','PF']
    param_list = [{"universe": '%s88' % c} for c in code_list]

    start = time.time()
    ret = year_split_file_tasks(__file__, param_list, __config__)
    print(json.dumps(ret, indent=2))
    print(time.time() - start)